import numpy as np
import cv2
from rknn.api import RKNN

RKNN_MODEL = "yoloV11_adblocker.rknn"
IMG_SIZE = 640
IMG_PATH = "./017.jpg"

if __name__ == "__main__":
    # 创建RKNN对象
    rknn = RKNN()

    # 从当前目录加载RKNN模型resnet_18
    ret = rknn.load_rknn(path=RKNN_MODEL)
    if ret != 0:
        print("Load Pytorch model failed!")
        exit(ret)

    # Init runtime environment
    print("--> Init runtime environment")
    ret = rknn.init_runtime(target="rk3566")
    if ret != 0:
        print("Init runtime environment failed!")
        exit(ret)
    print("done")

    # Set inputs
    img = cv2.imread(IMG_PATH)
    # img, ratio, (dw, dh) = letterbox(img, new_shape=(IMG_SIZE, IMG_SIZE))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))
    outputs = rknn.inference(inputs=[img], data_format="nhwc")
    print(outputs)
    # 所有待分类图片处理完后,释放RKNN对象
    rknn.release()
